﻿using Ndns.Packet.Abstractions;
using Ndns.Packet.Parser;
using Ndns.ServerHost.Abstractions;
using System;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace Ndns.Net
{
    public class DnsUdpServer : IDnsServer, IDisposable
    {

        public DnsProtocol Protocol => DnsProtocol.Udp;

        public EndPoint LocalEndPoint { get; set; }

        private volatile int _state;
        public DnsServerState State => (DnsServerState)_state;

        public object Tag { get; set; }

        public IDnsRequestHandler Handler { get; private set; }

        private bool _enableCompression = true;
        public bool EnableCompression
        {
            get => _enableCompression;
            set
            {
                if (_state == (int)DnsServerState.Running)
                    throw new InvalidOperationException();
                _enableCompression = value;
            }
        }

        public void SetHandler(IDnsRequestHandler handler)
        {
            if (handler == null)
                throw new ArgumentNullException(nameof(handler));
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpServer));
            if (_state == (int)DnsServerState.Running)
                throw new InvalidOperationException();

            Handler = handler;
        }

        private ReaderWriterLockSlim _rwLock = new ReaderWriterLockSlim();
        private volatile int _version;
        private Socket _socket;
        private readonly static byte[] s_bFalse = new byte[] { Convert.ToByte(false) };

        private DnsPacketParser _packetParser;

        public DnsUdpServer()
        {
            _packetParser = new DnsPacketParser();
        }
        public DnsUdpServer(DnsPacketParser packetParser)
        {
            _packetParser = packetParser ?? new DnsPacketParser();
        }
        public DnsUdpServer(IPEndPoint localEP)
        {
            LocalEndPoint = localEP ?? throw new ArgumentNullException(nameof(localEP));
            _packetParser = new DnsPacketParser();
        }
        public DnsUdpServer(IPEndPoint localEP, DnsPacketParser packetParser)
        {
            LocalEndPoint = localEP ?? throw new ArgumentNullException(nameof(localEP));
            _packetParser = packetParser ?? new DnsPacketParser();
        }

        public void Start()
        {
            //检查前置条件
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpServer));
            if (LocalEndPoint == null)
                throw new InvalidOperationException();
            if (Handler == null)
                throw new InvalidOperationException();
            if (_state == (int)DnsServerState.Running)
                throw new InvalidOperationException();

            //通知即将启动
            try
            {
                Handler.OnStarting(this);
            }
            catch (Exception)
            {
                _state = (int)DnsServerState.Error;
                return;
            }

            //新建Socket
            _socket = new Socket(LocalEndPoint.AddressFamily, SocketType.Dgram, ProtocolType.Udp);

            //设置Socket属性
            uint IOC_IN = 0x80000000;
            uint IOC_VENDOR = 0x18000000;
            uint SIO_UDP_CONNRESET = IOC_IN | IOC_VENDOR | 12;
            _socket.IOControl(unchecked((int)SIO_UDP_CONNRESET), s_bFalse, null);

            try
            {
                //绑定到端口
                _socket.Bind(LocalEndPoint);
            }
            catch (SocketException ex)
            {
                //设置状态并通知错误
                _state = (int)DnsServerState.Error;
                _socket.Dispose();
                try
                {
                    Handler.OnFatalError(this, ex);
                }
                catch (Exception)
                {
                    // Do nothing
                }
                return;
            }

            //设置状态并通知启动
            _state = (int)DnsServerState.Running;
            _version = unchecked(_version + 1);
            try
            {
                Handler.OnStarted(this);
            }
            catch (Exception)
            {
                // Do nothing
            }

            //开始接收
            Task.Run(() => SocketWorkerLoop(this));
        }

        public void Stop()
        {
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpServer));

            if (_state == (int)DnsServerState.Running)
            {
                try
                {
                    Handler.OnStopping(this);
                }
                catch (Exception)
                {
                    // Do nothing
                }

                _version = unchecked(_version + 1);
                
                _rwLock.EnterWriteLock();
                _socket.Dispose();
                _socket = null;
                _rwLock.ExitWriteLock();
                
            }
            _state = (int)DnsServerState.Stopped;

            try
            {
                Handler.OnStopped(this);
            }
            catch (Exception)
            {
                // Do nothing
            }
        }

        private static async Task SocketWorkerLoop(object state)
        {
            var server = (DnsUdpServer)state;
            var version = server._version;
            
            while (version == server._version)
            {
                var buffer = new byte[4096];
                SocketReceiveFromResult socketResult = default(SocketReceiveFromResult);
                try
                {
                    Task<SocketReceiveFromResult> taskReceiveFrom = null;
                    server._rwLock.EnterReadLock();
                    taskReceiveFrom = server._socket.ReceiveFromAsync(new ArraySegment<byte>(buffer, 0, 4096), SocketFlags.None, server.LocalEndPoint);
                    server._rwLock.ExitReadLock();
                    socketResult = await taskReceiveFrom.ConfigureAwait(false);
                }
                catch (ObjectDisposedException)
                {
                    break;
                }
                catch (Exception ex)
                {
                    Interlocked.CompareExchange(ref server._state, (int)DnsServerState.Error, (int)DnsServerState.Running);
                    server.Handler.OnFatalError(server, ex);
                    break;
                }
                StartForwardDnsRequest(server, version, buffer, socketResult, false);
            }
        }
        //把Task.Run放在单独的方法中以防止闭包导致的buffer和socketResult被共享的问题
        private static void StartForwardDnsRequest(DnsUdpServer server, int version, byte[] buffer, SocketReceiveFromResult socketResult, bool isBufferFromPool)
        {
            Task.Run(() => ForwardDnsRequest(server, version, socketResult.RemoteEndPoint, buffer, socketResult.ReceivedBytes, isBufferFromPool));
        }

        private static async Task ForwardDnsRequest(DnsUdpServer server, int version, EndPoint source, byte[] buffer, int length, bool isBufferFromPool)
        {
            //获得各个参数
            var enableCompression = server._enableCompression;

            //解析DNS数据包
            DnsPacket packet = null;
            using (var ms = new MemoryStream(buffer, 0, length, false))
            {
                server._packetParser.TryParse(ms, out packet, out var c);
            }

            //归还缓冲区
            if (isBufferFromPool)
            {
                //server._bufferPool.Return(buffer);
            }

            //检查解析结果
            if (packet == null)
                return;

            //上锁防止服务器被停止
            Task<DnsPacket> taskResponse = null;
            server._rwLock.EnterReadLock();
            try
            {
                //调用回调函数
                if (version == server._version)
                {
                    taskResponse = server.Handler.OnReceivedAsync(server, packet, source, CancellationToken.None);
                }
            }
            catch (Exception)
            {
                // Do nothing
            }
            finally
            {
                //解锁
                server._rwLock.ExitReadLock();
            }

            //等待回调函数结束
            DnsPacket response = null;
            try
            {
                if (taskResponse != null)
                    response = await taskResponse.ConfigureAwait(false);
            }
            catch (Exception)
            {
                // Do nothing
            }

            //检查回调函数结果
            if (response == null)
                return;

            //发送结果给客户端
            server._rwLock.EnterReadLock();
            try
            {
                Task<int> taskSendTo = null;
                if (version == server._version)
                {
                    taskSendTo = server._socket.SendToAsync(new ArraySegment<byte>(response.GetBytes(true, enableCompression)), SocketFlags.None, source);
                }
                server._rwLock.ExitReadLock();
                if (taskSendTo != null)
                    await taskSendTo.ConfigureAwait(false);
            }
            catch (ObjectDisposedException)
            {
                // Do nothing
            }
            catch (Exception)
            {
                // Do nothing
            }
        }

#region IDisposable Support
        private bool disposedValue = false; // 要检测冗余调用

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // 释放托管状态(托管对象)。
                    _socket?.Dispose();
                }

                // 释放未托管的资源(未托管的对象)并在以下内容中替代终结器。
                // 将大型字段设置为 null。
                _socket = null;
                _state = (int)DnsServerState.Stopped;

                disposedValue = true;
            }
        }

        // 添加此代码以正确实现可处置模式。
        public void Dispose()
        {
            // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
            Dispose(true);
        }

#endregion

    }
}
